#!/bin/bash

KUBE_AWS_CMD=${KUBE_AWS_CMD:-./bin/kube-aws}
E2E_DIR=$(cd $(dirname $0); pwd)
WORK_DIR=${E2E_DIR}/assets/${KUBE_AWS_CLUSTER_NAME}
TESTINFRA_DIR=${E2E_DIR}/testinfra
KUBE_AWS_TEST_INFRA_STACK_NAME=${KUBE_AWS_TEST_INFRA_STACK_NAME:-${KUBE_AWS_CLUSTER_NAME}-testinfra}
SRC_DIR=$(cd $(dirname $0); cd ..; pwd)
KUBECONFIG=${WORK_DIR}/kubeconfig
ETCD_COUNT=${ETCD_COUNT:-3}
CONTROLLER_COUNT=${CONTROLLER_COUNT:-2}
ETCD_VERSION=${ETCD_VERSION:-}

export KUBECONFIG

USAGE_EXAMPLE="KUBE_AWS_CLUSTER_NAME=kubeawstest1 KUBE_AWS_KEY_NAME=name/of/ec2/key KUBE_AWS_KMS_KEY_ARN=arn:aws:kms:us-west-1:<account id>:key/your-key KUBE_AWS_REGION=us-west-1 KUBE_AWS_AVAILABILITY_ZONE=us-west-1b ./$0 [init|configure|start|all]"

if [ "${KUBE_AWS_CLUSTER_NAME}" == "" ]; then
  echo KUBE_AWS_CLUSTER_NAME is not set. Run this command like $USAGE_EXAMPLE 1>&2
  exit 1
fi

if [ "${KUBE_AWS_KEY_NAME}" == "" ]; then
  echo KUBE_AWS_KEY_NAME is not set. Run this command like $USAGE_EXAMPLE 1>&2
  exit 1
fi

if [ "${KUBE_AWS_REGION}" == "" ]; then
  echo KUBE_AWS_REGION is not set. Run this command like $USAGE_EXAMPLE 1>&2
  exit 1
fi
# set the AWS CLI default region to the region we are deploying to, otherwise a mismatch with the user env will cause all `aws` commands to fail
export AWS_DEFAULT_REGION=${KUBE_AWS_REGION}

if [ "${KUBE_AWS_AVAILABILITY_ZONE}" == "" ]; then
  echo KUBE_AWS_AVAILABILITY_ZONE is not set. Run this command like $USAGE_EXAMPLE 1>&2
  exit 1
fi

if [ ! -e "${KUBE_AWS_CMD}" ]; then
  echo ${KUBE_AWS_CMD} does not exist. 1>&2
  exit 1
fi

KUBE_AWS_VERSION=$($KUBE_AWS_CMD version)
echo Using the kube-aws command at ${KUBE_AWS_CMD}"($KUBE_AWS_VERSION)". Set KUBE_AWS_CMD=path/to/kube-aws to override.

EXTERNAL_DNS_NAME=${KUBE_AWS_CLUSTER_NAME}.${KUBE_AWS_DOMAIN}
echo The kubernetes API would be accessible via ${EXTERNAL_DNS_NAME}

KUBE_AWS_S3_URI=${KUBE_AWS_S3_DIR_URI}/${KUBE_AWS_CLUSTER_NAME}
echo CloudFormation stack template would be uploaded to ${KUBE_AWS_S3_URI}

build() {
  echo Building kube-aws
  cd ${SRC_DIR}
  ./build
}

main_stack_name() {
  echo $KUBE_AWS_CLUSTER_NAME
}

main_status() {
  aws cloudformation describe-stacks --stack-name $(main_stack_name) --output json | jq -rc '.Stacks[0].StackStatus'
}

main_all() {
  status=$(aws cloudformation describe-stacks --stack-name $(main_stack_name) --output json | jq -rc '.Stacks[0].StackStatus')
  if [ "$status" = "" ]; then
    if [ ! -e "${WORK_DIR}/cluster.yaml" ]; then
      init
    fi
    if ! [ -d "${WORK_DIR}/credentials" -a -d "${WORK_DIR}/userdata" -a -e "${WORK_DIR}/kubeconfig" -a -d "${WORK_DIR}/stack-templates" ]; then
      configure
    fi
    if [ "$(main_status)" = "" ]; then
      up
    fi
    if [ "$KUBE_AWS_UPDATE" != "" ]; then
      scale_out
    fi
  fi
}

init() {
  echo Creating or ensuring existence of the kube-aws assets directory ${WORK_DIR}
  mkdir -p ${WORK_DIR}

  cd ${WORK_DIR}

  ${KUBE_AWS_CMD} init \
    --cluster-name ${KUBE_AWS_CLUSTER_NAME} \
    --external-dns-name ${EXTERNAL_DNS_NAME} \
    --region ${KUBE_AWS_REGION} \
    --availability-zone ${KUBE_AWS_AVAILABILITY_ZONE} \
    --key-name ${KUBE_AWS_KEY_NAME} \
    --kms-key-arn ${KUBE_AWS_KMS_KEY_ARN} \
    --hosted-zone-id ${KUBE_AWS_HOSTED_ZONE_ID}

  if [ "${KUBE_AWS_USE_CALICO}" != "" ]; then
    echo 'useCalico: true' >> cluster.yaml
  fi

  customize_cluster_yaml
}

regenerate_stack() {
  cd ${WORK_DIR}

  ${KUBE_AWS_CMD} render stack
}


regenerate_credentials() {
  cd ${WORK_DIR}

  rm -rf ./credentials

  ${KUBE_AWS_CMD} render credentials --generate-ca
}

configure() {
  cd ${WORK_DIR}

  rm -rf ./kubeconfig ./credentials ./userdata ./stack-templates ./exported

  ${KUBE_AWS_CMD} render stack
  ${KUBE_AWS_CMD} render credentials --generate-ca

  if [ "${EXISTING_CA}" != "" ]; then
    mv ./credentials/ca-key.pem my-ca-key.pem
    mv ./credentials/ca.pem my-ca.pem
    rm -rf ./credentials
    ${KUBE_AWS_CMD} render credentials --ca-key-path=./my-ca-key.pem --ca-cert-path=./my-ca.pem
  fi

  validate

  ${KUBE_AWS_CMD} up --export --s3-uri ${KUBE_AWS_S3_URI} --pretty-print

  echo Generated configuration files in ${WORK_DIR}:
  find .
}

kube-aws() {
  mkdir -p ${WORK_DIR}
  cd ${WORK_DIR}
  ${KUBE_AWS_CMD} "$@"
}

validate() {
  cd ${WORK_DIR}
  ${KUBE_AWS_CMD} validate --s3-uri ${KUBE_AWS_S3_URI}
}

customize_cluster_yaml() {
  echo Writing to $(pwd)/cluster.yaml

  if [ "${KUBE_AWS_DEPLOY_TO_EXISTING_VPC}" != "" ]; then
    echo -e "vpc:\n  id: $(testinfra_vpc)" >> cluster.yaml
    echo -e "routeTableId: $(testinfra_public_routetable)" >> cluster.yaml
    echo -e "workerSecurityGroupIds:\n- $(testinfra_glue_sg)" >> cluster.yaml
  fi

  if [ "${LIMIT_SSH_ACCESS}" != "" ]; then
    ip=$(curl -s https://api.ipify.org)
    echo -e "sshAccessAllowedSourceCIDRs:\n- $ip/32" >> cluster.yaml
  fi

  echo -e "
worker:
  nodePools:
  - name: asg1
    autoscaling:
      # this pool is an autoscaling-target of CA
      clusterAutoscaler:
        enabled: true
    # give this pool permissions to run CA
    clusterAutoscalerSupport:
      enabled: true
  - name: asg2
    autoScalingGroup:
      minSize: 0
      maxSize: 2
  - name: asg3
    count: 1
    waitSignal:
      enabled: true
    awsNodeLabels:
      enabled: true
    awsEnvironment:
      enabled: true
      environment:
        CFNSTACK: '{\"Ref\":\"AWS::StackId\"}'
    nodeDrainer:
      enabled: true
    nodeLabels:
      kube-aws.coreos.com/role: worker
      kube-aws.coreos.com/reservation-type: on-demand" >> cluster.yaml

  if [ "${KUBE_AWS_DEPLOY_TO_EXISTING_VPC}" != "" ]; then
    echo -e "
    securityGroupIds:
    - $(testinfra_glue_sg)
    loadBalancer:
      enabled: true
      names:
      - $(testinfra_public_elb)
      securityGroupIds:
      - $(testinfra_public_elb_backend_sg)
    targetGroup:
      enabled: true
      arns:
      - $(testinfra_target_group)
      securityGroupIds:
      - $(testinfra_public_alb_backend_sg)" >> cluster.yaml
  fi

  echo -e "
  - name: fleet1
    spotFleet:
      targetCapacity: 1
    clusterAutoscalerSupport:
      enabled: true
  - name: fleet2
    spotFleet:
      targetCapacity: 2
    awsNodeLabels:
      enabled: true
    awsEnvironment:
      enabled: true
      environment:
        CFNSTACK: '{\"Ref\":\"AWS::StackId\"}'
    autoscaling:
      clusterAutoscaler:
        enabled: true
    clusterAutoscalerSupport:
      enabled: true
    nodeDrainer:
      enabled: true
    nodeLabels:
      kube-aws.coreos.com/role: worker
      kube-aws.coreos.com/reservation-type: spot" >> cluster.yaml

  if [ "${KUBE_AWS_DEPLOY_TO_EXISTING_VPC}" != "" ]; then
    echo -e "
    securityGroupIds:
    - $(testinfra_glue_sg)
    loadBalancer:
      enabled: true
      names:
      - $(testinfra_public_elb)
      securityGroupIds:
      - $(testinfra_public_elb_backend_sg)
    targetGroup:
      enabled: true
      arns:
      - $(testinfra_target_group)
      securityGroupIds:
      - $(testinfra_public_alb_backend_sg)" >> cluster.yaml
  fi

  echo -e "
# controller configuration
controller:
  count: $CONTROLLER_COUNT
  nodeLabels:
    kube-aws.coreos.com/role: controller
waitSignal:
  enabled: true
kubeResourcesAutosave:
  enabled: true
experimental:
  awsNodeLabels:
    enabled: true
  auditLog:
    enabled: true" >> cluster.yaml

  if [ "${ENCRYPTION_AT_REST}" != "" ]; then
      echo -e "
  encryptionAtRest:
    enabled: true" >> cluster.yaml
  fi

  echo -e "
addons:
  clusterAutoscaler:
    enabled: true" >> cluster.yaml

  echo -e "
cloudWatchLogging:
  enabled: true
amazonSsmAgent:
  enabled: true
# etcd configuration
etcd:
  count: $ETCD_COUNT" >> cluster.yaml

  if [ "${ETCD_VERSION}" != "" ]; then
    echo -e "  version: ${ETCD_VERSION}" >> cluster.yaml
  fi

  if [ "${ETCD_DISASTER_RECOVERY_AUTOMATED}" != "" ]; then
    echo -e "  disasterRecovery:
    automated: true" >> cluster.yaml
  fi


  if [ "${ETCD_SNAPSHOT_AUTOMATED}" != "" ]; then
    echo -e "  snapshot:
    automated: true" >> cluster.yaml
  fi

  if [ "${ETCD_MEMBER_IDENTITY_PROVIDER}" != "" ]; then
    echo -e "  memberIdentityProvider: ${ETCD_MEMBER_IDENTITY_PROVIDER}" >> cluster.yaml
  fi

  if [ "${ETCD_INTERNAL_DOMAIN_NAME}" != "" ]; then
    echo -e "  internalDomainName: ${ETCD_INTERNAL_DOMAIN_NAME}" >> cluster.yaml
  fi
}

clean() {
  cd ${WORK_DIR}/..
  if [ -d "${KUBE_AWS_CLUSTER_NAME}" ]; then
    echo Removing the directory "${WORK_DIR}"
    rm -Rf ./${KUBE_AWS_CLUSTER_NAME}/*
  fi
}

up() {
  cd ${WORK_DIR}

  starttime=$(date +%s)

  ${KUBE_AWS_CMD} up --s3-uri ${KUBE_AWS_S3_URI} --pretty-print

  set +vx

  printf 'Waiting for the Kubernetes API to be accessible'

  while ! kubectl get no 2>/dev/null; do
    sleep 10
    printf '.'
  done

  endtime=$(date +%s)

  echo Done. it took $(($endtime - $starttime)) seconds until kubeapiserver to be in service.

  ${KUBE_AWS_CMD} status

  set -vx
}

status() {
  cd ${WORK_DIR}
  ${KUBE_AWS_CMD} status
}

main_destroy() {
  status=$(main_status)

  if [ "$status" != "" ]; then
    if [ "$status" != "DELETE_IN_PROGRESS" ]; then
        aws cloudformation delete-stack --stack-name $(main_stack_name)
        aws cloudformation wait stack-delete-complete --stack-name $(main_stack_name)
    else
        aws cloudformation wait stack-delete-complete --stack-name $(main_stack_name)
    fi
  else
    echo $(main_stack_name) does not exist. skipping.
  fi

  cd ${WORK_DIR}
  rm -Rf cluster.yaml user-data/ credentials/ stack-templates/
}

update() {
  cd ${WORK_DIR}

  ${KUBE_AWS_CMD} update --s3-uri ${KUBE_AWS_S3_URI} --pretty-print || true
  aws cloudformation wait stack-update-complete --stack-name ${KUBE_AWS_CLUSTER_NAME}
}

scale_out() {
  cd ${WORK_DIR}

  SED_CMD="sed -e 's/count: 1/count: 2/' -e 's/controllerCount: 2/controllerCount: 3/'"
  diff --unified cluster.yaml <(cat cluster.yaml | sh -c "${SED_CMD}") || true
  sh -c "${SED_CMD} -i bak cluster.yaml"
  ${KUBE_AWS_CMD} update --s3-uri ${KUBE_AWS_S3_URI}
  aws cloudformation wait stack-update-complete --stack-name ${KUBE_AWS_CLUSTER_NAME}

  printf 'Waiting for the Kubernetes API to be accessible'
  while ! kubectl get no 2>/dev/null; do
    sleep 10
    printf '.'
  done
  echo done
}

test-destruction() {
  aws cloudformation wait stack-delete-complete --stack-name ${KUBE_AWS_CLUSTER_NAME}
}

# Usage: DOCKER_REPO=quay.io/mumoshu/ SSH_PRIVATE_KEY=path/to/private/key ./e2e run conformance
conformance() {
  cd ${E2E_DIR}/kubernetes

  if [ "$DOCKER_REPO" == "" ]; then
    echo DOCKER_REPO is not set.
    exit 1
  fi

  if [ "$SSH_PRIVATE_KEY" == "" ]; then
    echo SSH_PRIVATE_KEY is not set.
    exit 1
  fi

  if [ ! -f "$SSH_PRIVATE_KEY" ]; then
    echo ${SSH_PRIVATE_KEY} does not exist.
    exit 1
  fi

  echo Opening ingress on 4194 and 10250...

  # Authorize these ingresses for E2E testing or it'll end up failing like:
  #
  # Summarizing 2 Failures:
  #
  # [Fail] [k8s.io] Proxy version v1 [It] should proxy logs on node using proxy subresource [Conformance]
  # /go/kubernetes/_output/local/go/src/k8s.io/kubernetes/test/e2e/proxy.go:325
  #
  # [Fail] [k8s.io] Proxy version v1 [It] should proxy logs on node with explicit kubelet port [Conformance]
  # /go/kubernetes/_output/local/go/src/k8s.io/kubernetes/test/e2e/proxy.go:325
  #
  # Ran 117 of 473 Specs in 3529.785 seconds
  # FAIL! -- 115 Passed | 2 Failed | 0 Pending | 356 Skipped --- FAIL: TestE2E (3529.90s)
  #
  group_id=$(aws cloudformation --output json describe-stack-resources --stack-name $(controlplane_stack_name) | jq -r '.StackResources[] | select(.LogicalResourceId == "SecurityGroupController").PhysicalResourceId')
  aws ec2 authorize-security-group-ingress --group-id ${group_id} --protocol tcp --port 4194 --source-group ${group_id} || echo 'skipping authorization for 4194'
  aws ec2 authorize-security-group-ingress --group-id ${group_id} --protocol tcp --port 10250 --source-group ${group_id} || echo 'skipping authorization for 10250'

  master_host=$(controller_host)

  echo Connecting to $master_host via SSH

  KUBE_AWS_ASSETS=${WORK_DIR} MASTER_HOST=$master_host make run-remotely
}

conformance_result() {
  cd ${E2E_DIR}/kubernetes

  master_host=$(controller_host)

  echo Connecting to $master_host via SSH

  KUBE_AWS_ASSETS=${WORK_DIR} MASTER_HOST=$master_host make show-log
}

controlplane_stack_name() {
  aws cloudformation --output json describe-stacks --stack-name ${KUBE_AWS_CLUSTER_NAME} | jq -r '.Stacks[].Outputs[] | select (.OutputKey == "ControlPlaneStackName").OutputValue'
}

first_host_for_asg() {
  controlplane_stack_name=$(controlplane_stack_name)
  aws ec2 describe-instances --output json --query "Reservations[].Instances[]
      | [?Tags[?Key==\`aws:cloudformation:stack-name\`].Value|[0]==\`${controlplane_stack_name}\`]
      | [?Tags[?Key==\`aws:cloudformation:logical-id\`].Value|[0]==\`${1}\`][]
      | [?State.Name==\`running\`][]
      | []" | jq -r 'map({InstanceId: .InstanceId, PublicIpAddress: .PublicIpAddress}) | first | .PublicIpAddress'
}

controller_host() {
  first_host_for_asg Controllers
}

worker_host() {
  first_host_for_asg Workers
}

ssh_controller() {
  ssh -o UserKnownHostsFile=/dev/null -o StrictHostKeyChecking=no -i ${KUBE_AWS_SSH_KEY} core@$(controller_host) "$@"
}

ssh_worker() {
  ssh -o UserKnownHostsFile=/dev/null -o StrictHostKeyChecking=no -i ${KUBE_AWS_SSH_KEY} core@$(worker_host) "$@"
}

testinfra_up() {
  cd ${TESTINFRA_DIR}

  aws cloudformation create-stack \
    --template-body file://$(pwd)/stack-template.yaml \
    --stack-name ${KUBE_AWS_TEST_INFRA_STACK_NAME} \
    --parameter ParameterKey=AZ1,ParameterValue=${KUBE_AWS_AZ_1} ParameterKey=Id,ParameterValue=${KUBE_AWS_CLUSTER_NAME}-infra
  aws cloudformation wait stack-create-complete \
    --stack-name ${KUBE_AWS_TEST_INFRA_STACK_NAME}
}

testinfra_update() {
  cd ${TESTINFRA_DIR}

  aws cloudformation update-stack \
    --template-body file://$(pwd)/stack-template.yaml \
    --stack-name ${KUBE_AWS_TEST_INFRA_STACK_NAME} \
    --parameter ParameterKey=AZ1,ParameterValue=${KUBE_AWS_AZ_1}
  aws cloudformation wait stack-update-complete \
    --stack-name ${KUBE_AWS_TEST_INFRA_STACK_NAME}
}

testinfra_destroy() {
  aws cloudformation delete-stack --stack-name ${KUBE_AWS_TEST_INFRA_STACK_NAME}
  aws cloudformation wait stack-delete-complete \
    --stack-name ${KUBE_AWS_TEST_INFRA_STACK_NAME}
}

testinfra_output() {
  aws cloudformation describe-stacks --stack-name ${KUBE_AWS_TEST_INFRA_STACK_NAME} | jq -r '.Stacks[0].Outputs[] | select(.OutputKey == "'$1'").OutputValue'
}

testinfra_vpc() {
  testinfra_output VPC
}

testinfra_public_routetable() {
  testinfra_output PublicRouteTable
}

testinfra_public_elb_backend_sg() {
  testinfra_output PublicELBBackendSG
}

testinfra_public_alb_backend_sg() {
  testinfra_output PublicALBBackendSG
}

testinfra_public_elb() {
  testinfra_output PublicELB
}

testinfra_target_group() {
  testinfra_output TargetGroup
}

testinfra_glue_sg() {
  testinfra_output GlueSG
}

all() {
  build

  if [ "${KUBE_AWS_DEPLOY_TO_EXISTING_VPC}" != "" ]; then
    testinfra_up
  fi

  main_all
  conformance
}

rerun() {
  all_destroy
  all
}

all_destroy() {
  main_destroy
  if [ "${KUBE_AWS_DEPLOY_TO_EXISTING_VPC}" != "" ]; then
    testinfra_destroy
  fi
}

kubesys() {
  kubectl --namespace kube-system "$@"
}

if [ "$1" == "" ]; then
  echo Usage: $USAGE_EXAMPLE
  exit 1
fi

set -vxe

"$@"
